# train_model.py
import pandas as pd
from sklearn.ensemble import IsolationForest
from sklearn.preprocessing import StandardScaler
import joblib
import os


def train_anomaly_detector(data_path="data/transaction_logs.csv", model_path="models/anomaly_model.pkl"):
    os.makedirs("models", exist_ok=True)

    print("🔍 加载数据...")
    df = pd.read_csv(data_path)
    X = df[['duration_sec', 'wait_time_ms', 'rows_affected', 'lock_count']]

    # 标准化
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    print("🧠 训练 Isolation Forest 模型...")
    model = IsolationForest(contamination=0.1, random_state=42, n_estimators=100)
    model.fit(X_scaled)

    # 保存模型和 scaler
    joblib.dump(model, model_path)
    joblib.dump(scaler, model_path.replace(".pkl", "_scaler.pkl"))

    print(f"✅ 模型已保存 → {model_path}")
    return model, scaler


if __name__ == "__main__":
    train_anomaly_detector()